/**
 * @file KDTree.cpp
 * @brief KD树的实现文件
 * 
 * 实现了高性能的KD树数据结构，特性包括：
 * 1. 内存对齐优化的节点分配
 * 2. 并行化的树构建
 * 3. 高效的最近邻搜索
 * 4. 八象限空间划分搜索
 */

#include "KDTree.h"
#include <algorithm>
#include <cmath>
#include <limits>
#include <memory>
#include <iostream>

/**
 * @brief 默认构造函数
 */
KDTree::KDTree() : root(nullptr) {}

/**
 * @brief 析构函数
 */
KDTree::~KDTree()
{
    cleanup(root);
}

/**
 * @brief 递归清理节点及其子树
 * @param node 要清理的节点指针
 */
void KDTree::cleanup(KDNode *node)
{
    if (!node)
        return;
    cleanup(node->left);
    cleanup(node->right);
    node->~KDNode();
    _aligned_free(node);
}

/**
 * @brief 构建KD树
 * @param points 输入点集，包含坐标和S1值的数组
 * @throw std::bad_alloc 当内存分配失败时抛出异常
 */
void KDTree::build(std::vector<std::pair<Point3D, double>> &points)
{
    #pragma omp parallel
    {
        #pragma omp single
        {
            root = buildTree(points, 0, 0, static_cast<int>(points.size()));
        }
    }
}

/**
 * @brief 递归构建KD树
 * @param points 输入点集，包含坐标和S1值
 * @param depth 当前递归深度
 * @param start 当前处理的点集起始索引
 * @param end 当前处理的点集结束索引
 * @return 返回构建的子树根节点
 * @throw std::bad_alloc 当内存分配失败时抛出异常
 */
KDTree::KDNode *KDTree::buildTree(std::vector<std::pair<Point3D, double>> &points,
                                 int depth, int start, int end)
{
    if (start >= end)
        return nullptr;

    // 选择分割轴
    int axis = depth % 3;
    int mid = (start + end) / 2;

    // 根据当前轴对点集进行分割
    std::nth_element(points.begin() + start, points.begin() + mid,
                    points.begin() + end,
                    [axis](const auto &a, const auto &b) {
                        if (axis == 0)
                            return a.first.x < b.first.x;
                        if (axis == 1)
                            return a.first.y < b.first.y;
                        return a.first.z < b.first.z;
                    });

    // 分配32字节对齐的节点内存
    void* ptr = _aligned_malloc(sizeof(KDNode), 32);
    if (!ptr) {
        throw std::bad_alloc();
    }

    // 构造新节点
    KDNode* node = ::new(ptr) KDNode(points[mid].first, points[mid].second);
    node->axis = axis;
    
    // 递归构建左右子树
    node->left = buildTree(points, depth + 1, start, mid);
    node->right = buildTree(points, depth + 1, mid + 1, end);

    return node;
}

/**
 * @brief 计算两点间的平方距离
 * @param p1 第一个点的坐标
 * @param p2 第二个点的坐标
 * @return 返回两点间距离的平方
 */
double KDTree::calculateDistanceSquared(const Point3D &p1, const Point3D &p2)
{
    return p1.distanceSquaredTo(p2);
}

/**
 * @brief 递归查找最近邻点
 * @param node 当前节点
 * @param target 目标点
 * @param nearest 当前找到的最近点集
 * @param max_dist 当前最大距离
 * @param k 需要查找的邻居数量
 */
void KDTree::findNearestNeighbors(KDNode *node, const Point3D &target,
                                 std::vector<std::pair<double, KDNode *>> &nearest,
                                 double &max_dist, int k) const
{
    if (!node) return;

    // 计算当前点到目标点的距离
    double dist = std::sqrt(calculateDistanceSquared(target, node->point));
    
    // 更新最近点集
    if (nearest.size() < k)
    {
        nearest.emplace_back(dist, node);
        std::push_heap(nearest.begin(), nearest.end());
        if (nearest.size() == k)
        {
            max_dist = nearest.front().first;
        }
    }
    else if (dist < max_dist)
    {
        std::pop_heap(nearest.begin(), nearest.end());
        nearest.back() = std::make_pair(dist, node);
        std::push_heap(nearest.begin(), nearest.end());
        max_dist = nearest.front().first;
    }

    // 确定搜索顺序
    int axis = node->axis;
    double diff = (axis == 0) ? target.x - node->point.x :
                 (axis == 1) ? target.y - node->point.y :
                              target.z - node->point.z;

    KDNode *first = (diff < 0) ? node->left : node->right;
    KDNode *second = (diff < 0) ? node->right : node->left;

    // 递归搜索子树
    findNearestNeighbors(first, target, nearest, max_dist, k);
    if (std::abs(diff) < max_dist)
    {
        findNearestNeighbors(second, target, nearest, max_dist, k);
    }
}

/**
 * @brief K近邻搜索的公共接口
 * @param target 目标点坐标
 * @param k 需要查找的邻居数量
 * @return 返回k个最近邻点及其距离，按距离升序排序
 */
std::vector<std::pair<double, KDTree::KDNode *>> KDTree::findKNearest(
    const Point3D &target, int k) const
{
    std::vector<std::pair<double, KDNode *>> nearest;
    double max_dist = std::numeric_limits<double>::max();

    findNearestNeighbors(root, target, nearest, max_dist, k);
    
    // 按距离排序结果
    std::sort(nearest.begin(), nearest.end(),
        [](const auto &a, const auto &b) {
            return a.first < b.first;
        });

    return nearest;
}

/**
 * @brief K近邻搜索的另一个重载版本
 * @param target 目标点坐标
 * @param k 需要查找的邻居数量
 * @param nearest 输出参数，存储查找到的k个最近邻点及其距离
 */
void KDTree::findKNearest(const Point3D &target, int k,
                         std::vector<std::pair<double, KDNode*>> &nearest) const
{
    nearest.clear();
    nearest.reserve(k);
    double max_dist = std::numeric_limits<double>::max();
    findNearestNeighbors(root, target, nearest, max_dist, k);
}

/**
 * @brief 在指定象限中递归查找最近点
 * @param node 当前节点
 * @param target 目标点坐标
 * @param nearest 当前找到的最近点集
 * @param max_dist 当前最大距离
 * @param k 需要查找的邻居数量
 * @param x_dir X方向限制(1=正,-1=负)
 * @param y_dir Y方向限制(1=正,-1=负)
 * @param z_dir Z方向限制(1=正,-1=负)
 */
void KDTree::findNearestNeighborsInOctant(KDNode *node, const Point3D &target,
                                         std::vector<std::pair<double, KDNode *>> &nearest,
                                         double &max_dist, int k,
                                         int x_dir, int y_dir, int z_dir) const
{
    if (!node)
        return;

    // 检查点是否在指定象限内
    bool in_octant = ((node->point.x - target.x) * x_dir >= 0) &&
                     ((node->point.y - target.y) * y_dir >= 0) &&
                     ((node->point.z - target.z) * z_dir >= 0);

    if (in_octant)
    {
        double dist_sq = calculateDistanceSquared(target, node->point);
        double dist = std::sqrt(dist_sq);

        // 更新最近点集
        if (nearest.size() < k)
        {
            nearest.push_back({dist, node});
            std::push_heap(nearest.begin(), nearest.end());
            if (nearest.size() == k)
            {
                max_dist = nearest.front().first;
            }
        }
        else if (dist < max_dist)
        {
            std::pop_heap(nearest.begin(), nearest.end());
            nearest.back() = {dist, node};
            std::push_heap(nearest.begin(), nearest.end());
            max_dist = nearest.front().first;
        }
    }

    // 确定搜索顺序
    int axis = node->axis;
    double diff = (axis == 0) ? target.x - node->point.x :
                 (axis == 1) ? target.y - node->point.y :
                              target.z - node->point.z;

    KDNode *first = (diff < 0) ? node->left : node->right;
    KDNode *second = (diff < 0) ? node->right : node->left;

    // 递归搜索子树
    findNearestNeighborsInOctant(first, target, nearest, max_dist, k, x_dir, y_dir, z_dir);
    if (diff * diff < max_dist * max_dist)
    {
        findNearestNeighborsInOctant(second, target, nearest, max_dist, k, x_dir, y_dir, z_dir);
    }
}

/**
 * @brief 在指定象限中查找近点
 * @param target 目标点坐标
 * @param x_dir X方向限制(1=正,-1=负)
 * @param y_dir Y方向限制(1=正,-1=负)
 * @param z_dir Z方向限制(1=正,-1=负)
 * @return 返回该象限中最近的点及其距离
 */
std::pair<double, KDTree::KDNode *> KDTree::findNearestInOctant(
    const Point3D &target, int x_dir, int y_dir, int z_dir) const
{
    std::vector<std::pair<double, KDNode *>> nearest;
    double max_dist = std::numeric_limits<double>::max();
    findNearestNeighborsInOctant(root, target, nearest, max_dist, 1, x_dir, y_dir, z_dir);

    if (nearest.empty())
    {
        return {std::numeric_limits<double>::max(), nullptr};
    }
    return nearest[0];
}

/**
 * @brief 在八个象限中查找最近点
 * @param target 目标点坐标
 * @return 返回八个象限中的最近点及其距离列表
 */
std::vector<std::pair<double, KDTree::KDNode *>> KDTree::findOctantNeighbors(const Point3D &target) const
{
    std::vector<std::pair<double, KDNode *>> octant_neighbors;
    octant_neighbors.reserve(8);

    // 定义八个象限的方向
    const int octants[8][3] = {
        {1, 1, 1},   // 右上前
        {1, 1, -1},  // 右上后
        {1, -1, 1},  // 右下前
        {1, -1, -1}, // 右下后
        {-1, 1, 1},  // 左上前
        {-1, 1, -1}, // 左上后
        {-1, -1, 1}, // 左下前
        {-1, -1, -1} // 左下后
    };

    // 在每个象限中搜索最近点
    for (int i = 0; i < 8; ++i)
    {
        auto nearest = findNearestInOctant(target,
                                         octants[i][0], octants[i][1], octants[i][2]);
        if (nearest.second != nullptr)
        {
            octant_neighbors.push_back(nearest);
        }
    }

    return octant_neighbors;
} 